package com.xkcyy.one.ai.agents;

import com.xkcyy.one.ai.common.ToolAgent;
import com.xkcyy.one.ai.common.tool.Tool;
import com.xkcyy.one.ai.llm.OpenAiClient;
import com.xkcyy.one.ai.llm.model.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 基于OpenAI的工具智能体实现
 * 支持工具调用和流式响应
 * 
 * @author AI Assistant
 */
@Slf4j
@Service
public class OpenAIToolAgent implements ToolAgent  {

    private final OpenAiClient openAiClient;
    private final Map<String, Tool> tools;
    
    @Autowired
    public OpenAIToolAgent(OpenAiClient openAiClient) {
        this.openAiClient = openAiClient;
        this.tools = new ConcurrentHashMap<>();
    }
    
    @Override
    public void registerTool(Tool toolExecutor) {
        tools.put(toolExecutor.getName(), toolExecutor);
        log.info("注册工具: {}", toolExecutor.getName());
    }
    
    @Override
    public void registerTools(List<Tool> executors) {
        executors.forEach(this::registerTool);
    }
    
//    @Override
//    public ChatResult chat(String prompt) {
//        ChatMessage userMessage = ChatMessage.builder()
//                .role("user")
//                .content(prompt)
//                .build();
//        return chat(Collections.singletonList(userMessage));
//    }
//
//    @Override
//    public ChatResult chat(List<ChatMessage> messages) {
//        // 准备工具列表
//        List<ToolDefinition> tools = prepareTools();
//        if (tools.isEmpty()) {
//            log.warn("没有可用的工具，将使用普通聊天");
//            return openAiClient.chat(messages);
//        }
//
//        // 使用工具进行聊天
//        ChatResult initialResponse = openAiClient.chatWithTools(messages, tools);
//
//        // 处理工具调用
//        if (initialResponse.getToolCalls() != null && !initialResponse.getToolCalls().isEmpty()) {
//            return handleToolCalls(messages, initialResponse);
//        }
//
//        return initialResponse;
//    }
    
    @Override
    public Flux<ChatResult> chatStream(String prompt) {
        ChatMessage userMessage = ChatMessage.builder()
                .role("user")
                .content(prompt)
                .build();
        return chatStream(Collections.singletonList(userMessage));
    }
    
    @Override
    public Flux<ChatResult> chatStream(List<ChatMessage> messages) {
        // 准备工具列表
        List<ToolDefinition> tools = prepareTools();
        if (tools.isEmpty()) {
            log.warn("没有可用的工具，将使用普通聊天流");
            return openAiClient.chatStream(messages);
        }

        // 用于收集工具调用和最终内容
        AtomicReference<String> finalContent = new AtomicReference<>("");
        List<ToolCallResult> collectedToolCalls = new ArrayList<>();
        // 使用带工具的流式聊天
        return openAiClient.chatStreamWithTools(messages, tools)
                .doOnNext(response -> {
                    // 累积内容
                    if (response.getContent() != null) {
                        finalContent.updateAndGet(current -> current + response.getContent());
                    }
                    
                    // 收集工具调用
                    if (response.getToolCalls() != null && !response.getToolCalls().isEmpty()) {
                        collectedToolCalls.addAll(response.getToolCalls());
                    }
                })
                .concatMap(response -> {
                    // 如果流还未结束，返回中间结果
                    if (!response.isComplete()) {
                        return Mono.just(response);
                    }

                    // 流结束时，如果有工具调用，执行它们
                    if (!collectedToolCalls.isEmpty()) {
                        List<ChatMessage> updatedMessages = new ArrayList<>(messages);
                        
                        // 添加助手消息
                        ChatMessage assistantMessage = ChatMessage.builder()
                                .role("assistant")
                                .content(finalContent.get())
                                .toolCalls(collectedToolCalls)
                                //.content(finalContent.get())
                                .build();
                        updatedMessages.add(assistantMessage);
                        
                        // 处理所有工具调用
                        for (ToolCallResult toolCall : collectedToolCalls) {
                            String toolName = toolCall.getFunction().getName();
                            Tool executor = this.tools.get(toolName);
                            
                            if (executor == null) {
                                log.error("找不到工具执行器: {}", toolName);
                                continue;
                            }
                            
                            // 执行工具
                            String toolResult = executor.execute(toolCall.getFunction().getArguments());
                            
                            // 添加工具响应消息
                            ChatMessage toolResponseMessage = ChatMessage.builder()
                                    .role("tool")
                                    .content(toolResult)
                                    .build();
                            updatedMessages.add(toolResponseMessage);
                        }
                        
                        // 继续流式对话
                        return this.chatStream(updatedMessages);
                    }
                    
                    // 如果没有工具调用，返回最终结果
                    return Mono.just(ChatResult.builder()
                            .content(finalContent.get())
                            .reasoningContent(response.getReasoningContent())
                            .toolCalls(collectedToolCalls)
                            .usage(response.getUsage())
                            .isComplete(true)
                            .build());
                });
    }
    
    /**
     * 准备可用的工具列表
     * 
     * @return 工具列表
     */
    private List<ToolDefinition> prepareTools() {
        if (tools.isEmpty()) {
            return Collections.emptyList();
        }
        
        return tools.values().stream()
                .map(Tool::getDefinition)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }
    
    /**
     * 处理非流式对话中的工具调用
     * 
     * @param messages 原始消息列表
     * @param response 包含工具调用的响应
     * @return 最终处理后的结果
     */
    private ChatResult handleToolCalls(List<ChatMessage> messages, ChatResult response) {
        List<ChatMessage> updatedMessages = new ArrayList<>(messages);
        
        // 添加助手消息
        ChatMessage assistantMessage = ChatMessage.builder()
                .role("assistant")
                .content(response.getContent())
                .build();
        updatedMessages.add(assistantMessage);
        
        // 依次处理每个工具调用并添加到消息列表
        for (ToolCallResult toolCall : response.getToolCalls()) {
            String toolName = toolCall.getFunction().getName();
            Tool executor = tools.get(toolName);
            
            if (executor == null) {
                log.error("找不到工具执行器: {}", toolName);
                continue;
            }
            
            // 执行工具
            String toolResult = executor.execute(toolCall.getFunction().getArguments());
            
            // 添加工具响应消息
            ChatMessage toolResponseMessage = ChatMessage.builder()
                    .role("tool")
                    .content(toolResult)
                    .build();
            updatedMessages.add(toolResponseMessage);
        }
        
        // 继续对话
        return openAiClient.chat(updatedMessages);
    }
    
    /**
     * 处理流式对话中的工具调用
     * 
     * @param messages 原始消息列表
     * @param response 包含工具调用的响应
     * @return 工具执行后继续的流
     */
    private Flux<ChatResult> handleToolCallsStream(List<ChatMessage> messages, ChatResult response) {
        List<ChatMessage> updatedMessages = new ArrayList<>(messages);
        
        // 添加助手消息
        ChatMessage assistantMessage = ChatMessage.builder()
                .role("assistant")
                .content(response.getContent())
                .build();
        updatedMessages.add(assistantMessage);
        
        // 处理多个工具调用
        List<ChatMessage> toolResponses = new ArrayList<>();
        for (ToolCallResult toolCall : response.getToolCalls()) {
            String toolName = toolCall.getFunction().getName();
            Tool executor = tools.get(toolName);
            
            if (executor == null) {
                log.error("找不到工具执行器: {}", toolName);
                continue;
            }
            
            // 执行工具
            String toolResult = executor.execute(toolCall.getFunction().getArguments());
            
            // 创建工具响应消息
            ChatMessage toolResponseMessage = ChatMessage.builder()
                    .role("tool")
                    .content(toolResult)
                    .build();
            toolResponses.add(toolResponseMessage);
        }
        
        // 添加所有工具响应
        updatedMessages.addAll(toolResponses);
        
        // 继续流式对话
        return this.chatStream(updatedMessages);
    }
} 